package bolts;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import backtype.storm.task.TopologyContext;
import backtype.storm.topology.BasicOutputCollector;
import backtype.storm.topology.OutputFieldsDeclarer;
import backtype.storm.tuple.Fields;
import backtype.storm.tuple.Tuple;
import backtype.storm.tuple.Values;

import common.TokenCount;

public class TokenFrequencyByUserBolt extends AbstractTimeWindowBolt {

    private Map<String, Map<String, TokenCount>> tokensByUser;

    private int userReservoirSize;

    private long currentTupleCount;

    public TokenFrequencyByUserBolt(String boltName, Map stormConf) {
        super(boltName, stormConf);
    }

    @Override
    public void prepare(Map stormConf, TopologyContext context) {
        super.prepare(stormConf, context);

        this.userReservoirSize = Integer.valueOf((String) stormConf.get("userReservoirSize"));
        this.tokensByUser = new HashMap<String, Map<String, TokenCount>>();
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        declarer.declare(new Fields("userTokenCounts"));
    }

    @Override
    protected void handleNormalTuple(Tuple input, BasicOutputCollector collector) {
        try {
            currentTupleCount++;

            final String user = input.getStringByField("screenName");
            boolean shouldSample = false;
            if (!tokensByUser.containsKey(user)) {
                // should sample or not
                if (tokensByUser.size() < userReservoirSize) {
                    shouldSample = true;
                } else if (Math.random() <= ((double) userReservoirSize / (double) currentTupleCount)) {
                    shouldSample = true;
                } else {
                    shouldSample = false;
                }
            } else {
                shouldSample = true;
            }

            if (shouldSample) {
                // if reservoir is full, remove random item
                if (tokensByUser.size() == userReservoirSize) {
                    final int indexToRemove = (int) (Math.random() * userReservoirSize);
                    final String userToRemove = (String) tokensByUser.keySet().toArray()[indexToRemove];
                    tokensByUser.remove(userToRemove);
                }

                if (!tokensByUser.containsKey(user)) {
                    tokensByUser.put(user, new HashMap<String, TokenCount>());
                }

                final Map<String, TokenCount> userTokens = tokensByUser.get(user);
                final String[] tokens = input.getStringByField("text").toLowerCase().split(" ");
                for (final String token : tokens) {
                    if (userTokens.containsKey(token)) {
                        userTokens.get(token).incrementCount();
                    } else {
                        userTokens.put(token, new TokenCount(token));
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    protected void handleTickTuple(Tuple input, BasicOutputCollector collector) {
        final Map<String, List<TokenCount>> output = new HashMap<String, List<TokenCount>>();

        for (final Entry<String, Map<String, TokenCount>> entry : tokensByUser.entrySet()) {
            final String user = entry.getKey();

            final List<TokenCount> tokenCounts = new ArrayList<TokenCount>(entry.getValue().size());
            for (final Entry<String, TokenCount> countEntry : entry.getValue().entrySet()) {
                tokenCounts.add(countEntry.getValue());
            }
            output.put(user, tokenCounts);
        }
        tokensByUser.clear();
        currentTupleCount = 0;

        collector.emit(new Values(output));
    }

}
